# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Interaction layers
"""

from typing import Union, Tuple, Optional

import torch
from torch import nn, Tensor

from ..utils import get_integer, GLOBAL_DEVICE
from ..activation import get_activation

_INTERACTION_BY_KEY ={}


def _interaction_register(*aliases):
    """Return the alias register."""
    def alias_reg(cls):
        name = cls.__name__.lower()
        if name not in _INTERACTION_BY_KEY:
            _INTERACTION_BY_KEY[name] = cls

        for alias in aliases:
            if alias not in _INTERACTION_BY_KEY:
                _INTERACTION_BY_KEY[alias] = cls

        return cls
    
    return alias_reg


# Define the Interaction class
class Interaction(nn.Module):
    """
    Interaction layer network
    
    Args:
        dim_node_rep (int): Dimension of node representations.
        dim_edge_rep (int): Dimension of edge representations.
        dim_node_emb (int, optional): Dimension of node embeddings. Defaults to None.
        dim_edge_emb (int, optional): Dimension of edge embeddings. Defaults to None.
        activation (str, optional): Name of the activation function. Defaults to 'silu'.
    """
    def __init__(self,
                 dim_node_rep: int,
                 dim_edge_rep: int,
                 dim_node_emb: int = None,
                 dim_edge_emb: int = None,
                 activation: Union[nn.Module, str] = 'silu',
                 **kwargs,
                 ):
        super().__init__()
        self._kwargs = kwargs

        self.dim_node_rep = get_integer(dim_node_rep)
        self.dim_edge_rep = get_integer(dim_edge_rep)
        self.dim_node_emb = get_integer(dim_node_emb)
        self.dim_edge_emb = get_integer(dim_edge_emb)
        if self.dim_edge_emb is None:
            self.dim_edge_emb = self.dim_node_rep
        self.activation = get_activation(activation)

        self.device = GLOBAL_DEVICE()

    def _print_info(self,
                    num_retraction: int = 6,
                    num_gap: int = 3,
                    char: str = '-'
                    ) -> None:
        ret = char * num_retraction
        gap = char * num_gap
        print(ret + gap + f' Feature dimension: {self.dim_node_rep}')
        print(ret + gap + f' Activation function: {self.activation.cls_name}')

    def forward(self, 
                node_vec: Tensor, 
                node_emb: Tensor, 
                node_mask: Tensor, 
                edge_vec: Tensor, 
                edge_emb: Tensor, 
                edge_mask: Optional[Tensor] = None, 
                edge_cutoff: Optional[Tensor] = None,
                neigh_idx: Optional[Tensor] = None,
                **kwargs
                ) -> Tuple[Tensor, Tensor]:
        """
        Forward pass of the interaction layer.
        
        Args:
            node_vec: Tensor of node vectors.
            node_emb: Tensor of node embeddings.
            node_mask: Tensor indicating valid nodes.
            edge_vec: Tensor of edge vectors.
            edge_emb: Tensor of edge embeddings.
            edge_mask: Optional tensor indicating valid edges.
            edge_cutoff: Optional tensor for edge cutoff values.
            neigh_idx: Optional tensor of neighbor indices.
        
        Returns:
            node_vec: Tensor of updated node vectors.
            edge_vec: Tensor of updated edge vectors.
        """
        return node_vec, edge_vec